2.1 Тест-кейсы¶

В процессе этого задания, я протестирую финальный алгоритм на 7 тестовых наборах, и оценю результат

Импортирование библиотек¶

In [11]:
# библиотека с предобученной моделью YOLO
from ultralytics import YOLO

# библиотека для работы с регулярными выражениями
import re

# библиотеки для работы с изображениями
import cv2
from PIL import Image
import matplotlib.pyplot as plt

# библиотека с OCR моделью
import easyocr

Инициализирую все нужные модели

In [12]:
# модель для детекции ценника из прошлого модуля
model = YOLO(r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\runs\detect\price_detection_v42\weights\best.pt')

# OCR модель для распознавания цены
ocr = easyocr.Reader(['en'])

# предопределенная картинка для функции
img = r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\31.jpg'
Neither CUDA nor MPS are available - defaulting to CPU. Note: This module is much faster with a GPU.
In [13]:
def rec_price(det_model: YOLO = model, 
              ocr: easyocr.Reader = ocr, 
              img_dir: str = img):
    
    '''
        Процедура для обнаружения и распознавания цен на изображении.

    Args:
        - det_model (YOLO): Модель YOLO для обнаружения bounding box'ов цен.
        - ocr (easyocr.Reader): Модель OCR для распознавания текста.
        - img_dir (str): Путь к изображению, на котором нужно найти цену.

    Returns:
        list: Функция отображает изображение с обнаруженными ценами и возвращает список с определенными ценами
    '''    

    image = cv2.imread(img_dir)
    res = det_model(image)
    
    image_ocr = cv2.imread(img_dir, cv2.IMREAD_GRAYSCALE)
    _, binary_image = cv2.threshold(image_ocr, 100, 255, cv2.THRESH_BINARY)


    image_ocr = cv2.equalizeHist(binary_image)

    prices = []

    for result in res:
        boxes = result.boxes.xyxy.cpu().numpy()
        for box in boxes:
            x1, y1, x2, y2 = map(int, box)
            correct = (x2-x1)*0.28
            print(correct)
            x2 = int(x2-correct)
            crop = image_ocr[y1:y2, x1:x2]

            # --------------------------------------------

            # plt.figure(figsize=(5, 5))
            # plt.imshow(crop, cmap='gray')
            # plt.title("Cropped Image")               # отображение обрезанной части картинки
            # plt.axis('off')
            # plt.show()

            # --------------------------------------------

            ocr_res = ocr.readtext(crop, allowlist='0123456789') # допускаем к распознаванию только цифры

            price = None 

            for detection in ocr_res:
                price = detection[1]
                confidence = detection[2]

                match = re.search(r'\d+[\.,]?\d*', price)  # на всякий случай проверяем регуляркой

                if match:
                    price = match.group()
                    print(f'Price: {price}\nConfidence: {confidence:.2f}')
                    break
            if price is not None:
                cv2.rectangle(image, (x1,y1), (x2+int(correct),y2), (255, 97, 0), 2)  # выделение изначального бокса
                cv2.rectangle(image, (x1,y1), (x2,y2), (97, 255, 0), 2)     # выделение инпута в ocr
                cv2.putText(image, price + 'rub', (x1,y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 97, 0), 2)
                prices.append(price)

    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(img_rgb)
    pil_img.show() 

    return prices

1. Нарисованные ценники¶

In [14]:
# список с путями к изображениям
images = [
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\11.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\12.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\13.jpg'
          ]

# проходимся по каждой фотографии
for image in images:
    rec_price(det_model=model, ocr=ocr, img_dir=image)
    print('--------------------------------------')
0: 288x640 (no detections), 83.3ms
Speed: 4.0ms preprocess, 83.3ms inference, 1.0ms postprocess per image at shape (1, 3, 288, 640)
--------------------------------------

0: 288x640 (no detections), 84.9ms
Speed: 2.1ms preprocess, 84.9ms inference, 3.8ms postprocess per image at shape (1, 3, 288, 640)
--------------------------------------

0: 288x640 (no detections), 105.1ms
Speed: 5.0ms preprocess, 105.1ms inference, 1.6ms postprocess per image at shape (1, 3, 288, 640)
--------------------------------------

Результат: ни распознал ни одну

2. Paint-ценники¶

In [17]:
# список с путями к изображениям
images = [
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\15.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\16.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\17.jpg',
          ]

# проходимся по каждой фотографии
for image in images:
    rec_price(det_model=YOLO(r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\runs\detect\price_detection_v3\weights\best.pt'), ocr=ocr, img_dir=image)
    print('--------------------------------------')
0: 352x640 1 price, 62.4ms
Speed: 1.0ms preprocess, 62.4ms inference, 2.0ms postprocess per image at shape (1, 3, 352, 640)
35.0
Price: 682
Confidence: 0.57
--------------------------------------

0: 384x640 1 price, 77.3ms
Speed: 7.3ms preprocess, 77.3ms inference, 1.0ms postprocess per image at shape (1, 3, 384, 640)
34.720000000000006
Price: 688
Confidence: 1.00
--------------------------------------

0: 320x640 (no detections), 124.6ms
Speed: 3.7ms preprocess, 124.6ms inference, 0.7ms postprocess per image at shape (1, 3, 320, 640)
--------------------------------------

Результат: распознал 2 из 3

на старой версии модели

image.png

3. Ценники из буфета¶

In [26]:
# список с путями к изображениям
images = [
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\21.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\22.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\23.jpg',
          ]

# проходимся по каждой фотографии
for image in images:
    rec_price(det_model=model, img_dir=image)
    print('--------------------------------------')
0: 288x640 (no detections), 87.4ms
Speed: 4.6ms preprocess, 87.4ms inference, 0.0ms postprocess per image at shape (1, 3, 288, 640)
--------------------------------------

0: 320x640 (no detections), 75.4ms
Speed: 3.2ms preprocess, 75.4ms inference, 1.0ms postprocess per image at shape (1, 3, 320, 640)
--------------------------------------

Результат: ни распознал ни одну

4. Ценники Пятерочки из Интернета¶

In [5]:
# список с путями к изображениям
images = [
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\41.jfif',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\42.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\41.jpg',
          ]

# проходимся по каждой фотографии
for image in images:
    rec_price(det_model=model, img_dir=image)
    print('--------------------------------------')
0: 480x640 5 prices, 108.8ms
Speed: 7.9ms preprocess, 108.8ms inference, 1.2ms postprocess per image at shape (1, 3, 480, 640)
63.28000000000001
Price: 48
Confidence: 1.00
85.12
Price: 16
Confidence: 1.00
63.56000000000001
Price: 44
Confidence: 1.00
60.2
Price: 16
Confidence: 1.00
71.68
Price: 34
Confidence: 0.69
--------------------------------------

0: 480x640 3 prices, 74.8ms
Speed: 2.0ms preprocess, 74.8ms inference, 1.1ms postprocess per image at shape (1, 3, 480, 640)
52.080000000000005
Price: 179
Confidence: 0.58
36.400000000000006
Price: 119
Confidence: 0.99
14.560000000000002
--------------------------------------

0: 480x640 2 prices, 109.0ms
Speed: 3.9ms preprocess, 109.0ms inference, 2.0ms postprocess per image at shape (1, 3, 480, 640)
134.96
Price: 120
Confidence: 1.00
110.04
Price: 115
Confidence: 1.00
--------------------------------------

Результат: определил 3 из 3

image.png image-2.png image-3.png

5. Ценники Магнита из Интернета¶

In [19]:
# список с путями к изображениям
images = [
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\51.jfif',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\52.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\53.jpg',
          ]

# проходимся по каждой фотографии
for image in images:
    rec_price(det_model=model, img_dir=image)
    print('--------------------------------------')

0: 640x480 2 prices, 86.4ms
Speed: 4.4ms preprocess, 86.4ms inference, 1.0ms postprocess per image at shape (1, 3, 640, 480)
37.52
Price: 913
Confidence: 0.89
44.24
Price: 109
Confidence: 1.00
--------------------------------------

0: 480x640 2 prices, 90.9ms
Speed: 3.0ms preprocess, 90.9ms inference, 1.7ms postprocess per image at shape (1, 3, 480, 640)
36.96
Price: 541
Confidence: 0.69
35.0
Price: 99
Confidence: 1.00
--------------------------------------

0: 480x640 (no detections), 214.1ms
Speed: 16.2ms preprocess, 214.1ms inference, 0.5ms postprocess per image at shape (1, 3, 480, 640)
--------------------------------------

Результат: определил 3 из 3

image.png image-2.png

6. Ценники из тестового набора ч.1¶

In [ ]:
# список с путями к изображениям
images = [
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\61.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\62.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\63.jpg',
          ]

# проходимся по каждой фотографии
for image in images:
    rec_price(det_model=model, img_dir=image)
    print('--------------------------------------')

0: 480x640 1 price, 86.1ms
Speed: 4.0ms preprocess, 86.1ms inference, 2.0ms postprocess per image at shape (1, 3, 480, 640)
71.4
Price: 29
Confidence: 0.67
--------------------------------------

0: 480x640 1 price, 70.8ms
Speed: 1.9ms preprocess, 70.8ms inference, 1.1ms postprocess per image at shape (1, 3, 480, 640)
54.32000000000001
--------------------------------------

0: 480x640 2 prices, 123.4ms
Speed: 2.3ms preprocess, 123.4ms inference, 2.5ms postprocess per image at shape (1, 3, 480, 640)
46.2
Price: 499
Confidence: 0.98
39.480000000000004
Price: 319
Confidence: 1.00
--------------------------------------

Результат: определил 2 из 3

image.png image-2.png image-3.png

7. Ценники из тестового набора ч.2¶

In [20]:
# список с путями к изображениям
images = [
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\77.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\71.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\72.jpg',
    r'D:\Helper\MLBazyak\homework\06_01\06_01_hw\hih\73.jpg',
          ]

# проходимся по каждой фотографии
for image in images:
    rec_price(det_model=model, img_dir=image)
    print('--------------------------------------')
0: 640x480 (no detections), 130.3ms
Speed: 16.0ms preprocess, 130.3ms inference, 2.6ms postprocess per image at shape (1, 3, 640, 480)
--------------------------------------

0: 480x640 1 price, 85.5ms
Speed: 7.2ms preprocess, 85.5ms inference, 4.3ms postprocess per image at shape (1, 3, 480, 640)
37.52
--------------------------------------

0: 480x640 1 price, 212.1ms
Speed: 5.4ms preprocess, 212.1ms inference, 7.2ms postprocess per image at shape (1, 3, 480, 640)
99.68
Price: 279
Confidence: 0.98
--------------------------------------

0: 480x640 3 prices, 710.1ms
Speed: 17.0ms preprocess, 710.1ms inference, 2.0ms postprocess per image at shape (1, 3, 480, 640)
47.040000000000006
Price: 32
Confidence: 0.99
43.68000000000001
Price: 32
Confidence: 0.99
32.760000000000005
Price: 82
Confidence: 0.77
--------------------------------------

Результат: 2 из 3

image-4.png image-5.png image.png

Рефлексия¶

  • результат выше среднего, но модели с большим кол-вом эпох хуже определяют различные ценники под углом или наклоном
  • можно попробовать объединить 3, 10 и 20 эпошные модели в ансамбль, и определять лучший результат голосованием